from leetcode_test import TreeNode

#　大神的代码，根本写不出来啊

class Solution(object):
    def deleteNode(self, root, key):
        """
        :type root: TreeNode
        :type key: int
        :rtype: TreeNode
        """

        if not root:
            return None
        if root.val > key:
            root.left = self.deleteNode(root.left, key)
        elif root.val < key:
            root.right = self.deleteNode(root.right, key)
        else:
            if not root.left or not root.right:
                root = root.left if root.left else root.right
            else:
                cur = root.right
                while cur.left:
                    cur = cur.left
                root.val = cur.val
                root.right = self.deleteNode(root.right, cur.val)

        return root


def test():
    t = TreeNode.create([5, 3, 6, 2, 4, None, 7])
    s = Solution().deleteNode(t, 3)
    assert s.get_value() == [5, 4, 6, 2, None, None, 7]


def test1():
    t = TreeNode.create([5, 3, 6, 2, 4])
    s = Solution().deleteNode(t, 6)
    assert s.get_value() == [5, 3, None, 2, 4]

def test2():
    t = TreeNode.create([5, 3, 6, 2, 4])
    s = Solution().deleteNode(t, 5)
    assert s.get_value() == [6,3 ,None, 2, 4]

if __name__ == "__main__":
    
    t = TreeNode.create([5, 3, 6, 2, 4])
    s = Solution().deleteNode(t, 5)
    print(t.get_value())
